import numpy as np
import statsmodels.formula.api as smf
from sklearn.model_selection import train_test_split
from sqlalchemy import create_engine
import pymysql
from sklearn.metrics import r2_score, mean_absolute_error
import pandas as pd
import matplotlib.pyplot as plt

#数据库连接配置 数据库地址、用户名、密码、数据库名、端口
db_conf = {
    'host':'localhost',
    'user':'root',
    'password':'MYSQL123',
    'database':'mysql',
    'port':3306,
    'charset':'utf8mb4',
}
engire = create_engine(
    f"mysql+pymysql://{db_conf['user']}:{db_conf['password']}@{db_conf['host']}:{db_conf['port']}/{db_conf['database']}")
conn = pymysql.connect(**db_conf)
chunk_size = 10000
#筛选银行日线数据。 添加模拟股票代码
# SELECT d.* FROM date_1 d WHERE d.trade_date BETWEEN '2023-01-01' and '2023-12-31' and d.ts_code = '000001.SZ'
df = pd.read_sql_query(""" 
    SELECT d.*, m.net_mf_vol, m.sell_elg_vol, m.buy_elg_vol, m.sell_lg_vol,m.buy_lg_vol,i.closes,i.vol AS i_vol
    FROM date_1 d
    JOIN moneyFlows m
    ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
    LEFT JOIN index_daily i ON d.trade_date = i.trade_date and i.ts_code = '000001.SH'
    WHERE d.trade_date BETWEEN '2023-01-01' and '2023-12-31' and d.ts_code = '000001.SZ'
    """,
    conn,
    chunksize = chunk_size)
df1 = pd.concat(df,ignore_index=True)
print(df1.head)
#增加一列，股票的涨跌幅
# df1['closes'] = pd.to_numeric(df1['closes'], errors='coerce')
# df1['zd_close'] = df1['closes'].pct_change().round(2).fillna(0)

# df1['i_vol'] = pd.to_numeric(df1['i_vol'], errors='coerce')
# df1['zd_closes'] = df1['i_vol'].pct_change().round(2).fillna(0)
df1['zd_closes'] = round((df1['i_vol'] - df1['i_vol'].shift(1)) / df1['i_vol'].shift(1), 2)  # 分母应为 shift(1)
df1['zd_close'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2 )
df1['hz_close'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1), 2 )
df1['hz_vol'] = round((df1['i_vol'] - df1['i_vol'].shift(1)) / df1['i_vol'].shift(2), 2 )
print(df1.head)

#处理缺失值
df1 = df1.dropna(subset = ['zd_close','hz_close','hz_vol'])
numeric_cols = df1.select_dtypes(include = ['number']).columns.tolist()
# #处理缺少数据
# df1 = df1.dropna(subset = ['zd_closes'])
# print(df1.head)
# ex = ['id','ts_code','trade_date','the_date','opens','high','low','closes','pre_closes','changes','pct_chg']

#筛选合适的自变量
# number = df1.select_dtypes(include = ['number']).columns.tolist()
# newList = [col for col in number if col not in ex]
excluded = ['zd_close','ts_code','trade_date','id','pre_closes','changes','pct_chg','opens','high','low','closes',
            'vol','buy_elg_vol','i_closes','i_vol','sell_lg_vol']#根据实际列名调整
predictors = [col for col in numeric_cols if col not in excluded]
#print(newList)
formula = 'zd_closes ~ ' + ' + '.join(predictors)
results1 = smf.ols(formula,data=df1).fit()
# formuls = 'zd_closes ~ ' + ' + '.join(newList)
res = smf.ols(formula, data=df1).fit()
print(results1.summary())




plt.figure(figsize = (10, 6))
plt.scatter(results1.fittedvalues, df1['zd_closes'])
plt.title('Fitted Values vs. Observed Values')
plt.xlabel('Fitted Values')
plt.ylabel('Observed Values')
plt.show()


features = df1[['amoint','net_mf_vol','sell_elg_vol','buy_lg_vol','hz_close','hz_vol',]]
correlation_matrix = features.corr()

#Plotting a heatmap for the correlation matriix
plt.figure(figsize = (10, 6))
plt.heatmap(correlation_matrix,annot = True, cmap = 'coolwarm', cbar = True, fmt = '.2f', linewidths = .5)
plt.title('Correlation Matrix between Features')
plt.show()

# # 关闭数据库连接
# conn.close()
# engine.dispose()